package com.zhiyou.app.service.impl;

import cn.hutool.core.io.FileUtil;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.xinchao.components.redis.utils.RedisUtil;
import com.zhiyou.app.common.exception.BizException;
import com.zhiyou.app.common.ws.DefaultTaskFinishHandler;
import com.zhiyou.app.common.ws.TaskProcessTracker;
import com.zhiyou.app.config.oss.OssUtil;
import com.zhiyou.app.pojo.entity.TaskEntity;
import com.zhiyou.app.pojo.enums.TaskStatusEnum;
import com.zhiyou.app.service.ComfyUIService;
import com.zhiyou.app.service.ITaskService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static com.zhiyou.app.common.constant.CommonConstant.REDIS_TASK_ID_PREFIX;

@Slf4j
@Service
public class ComfyUIServiceImpl implements ComfyUIService {
    @Value("${comfyui.server}")
    private String comfyUiServer;
    @Resource
    private ITaskService taskService;
    @Resource
    private RedisUtil redisUtil;
    @Resource
    private OssUtil ossUtil;
    @Resource
    private TaskProcessTracker taskProcessTracker;
    @Resource
    private DefaultTaskFinishHandler defaultTaskFinishHandler;

    @Override
    public String uploadFile(String fileUrl) {
        File tempFile = null;
        try {
            File currentDir = new File(".");
            String currentDirPath = currentDir.getCanonicalPath();
            if (!fileUrl.contains("/")) {
                throw new BizException("原始文件路径异常");
            }
            tempFile = new File(currentDirPath + fileUrl.substring(fileUrl.lastIndexOf("/")));
            try {
                HttpUtil.downloadFile(fileUrl, tempFile);
            } catch (Exception e) {
                log.error("文件下载失败", e);
                throw new BizException("下载文件失败");
            }
            HttpResponse response = HttpUtil.createPost(comfyUiServer + "upload/image").header("Content-Type", "multipart/form-data").form("image", tempFile).execute();
            if (response.getStatus() != 500) {
                JSONObject jsonObject = JSONUtil.parseObj(response.body());
                log.info(jsonObject.toString());
                return jsonObject.getStr("name");
            } else {
                throw new RuntimeException(response.body());
            }
        } catch (RuntimeException | IOException e) {
            throw new RuntimeException(e);
        } finally {
            if (tempFile != null && tempFile.exists()) {
                FileUtil.del(tempFile);
            }
        }
    }

    @Override
    public Map<String, List<String>> getRes(String promptId, TaskEntity taskEntity) {
        Map<String, List<String>> images = new HashMap<>();
        HttpResponse response = HttpUtil.createGet(comfyUiServer + "history/" + promptId).execute();
        if (response.getStatus() == 200) {
            JSONObject jsonObject = JSONUtil.parseObj(response.body());
            if (!jsonObject.toString().equals("{}")) {
                for (Map.Entry<String, Object> outputs : jsonObject.getJSONObject(promptId).getJSONObject("outputs")) {
                    JSONArray imagesArr = new JSONArray();
                    if (outputs.getValue() instanceof JSONObject) {
                        imagesArr = ((JSONObject) outputs.getValue()).getJSONArray("images");
                    }
                    if (imagesArr != null) {
                        List<String> imgList = new ArrayList<>();
                        List<JSONObject> tempList = imagesArr.stream().map(i -> (JSONObject) i).filter(i -> i.getStr("type").equals("output")).collect(Collectors.toList());
                        for (JSONObject entries : tempList) {
                            imgList.add(view(entries, taskEntity));
                        }
                        images.put(outputs.getKey(), imgList);
                    }
                }
            } else {
                log.debug("任务{}正在执行或执行失败", promptId);
            }
        }
        return images;
    }

    @Override
    public String view(JSONObject res, TaskEntity taskEntity) {
        try {
            return ossUtil.uploadAndUseFile(taskEntity.getOrderNum() + "_" + ".png", "idphoto",
                    new URL(comfyUiServer + "view?filename=" + res.getStr("filename") + "&subfolder=" + res.getStr("subfolder") + "&type=" + res.getStr("type")).openStream());
        } catch (IOException e) {
            log.error("文件流读取失败", e);
            throw new BizException("文件流读取失败", e.getMessage());
        }
    }


    @Override
    public void doTask(String promptJson, TaskEntity taskEntity, String clientId) {
        HttpResponse response = HttpUtil.createPost(comfyUiServer + "prompt").header("Content-Type", "application/json").body(promptJson).execute();
        if (response.getStatus() == 200) {
            JSONObject responseObj = JSONUtil.parseObj(response.body());
            log.info("任务提交参数：{}", responseObj);
            taskEntity.setTaskStatus(TaskStatusEnum.PROCESSING);
            taskEntity.setTaskPid(responseObj.getStr("prompt_id"));
            taskService.updateById(taskEntity);
            redisUtil.setEx(String.format(REDIS_TASK_ID_PREFIX, taskEntity.getTaskPid()), taskEntity.getTaskType(), 10, TimeUnit.DAYS);
            log.info("comfyui提交任务成功：{}", clientId);
        } else {
            log.error("任务提交失败:{}", response.body());
            taskEntity.setTaskStatus(TaskStatusEnum.ERROR);
            taskService.updateById(taskEntity);
            throw new BizException("任务提交失败");
        }
    }

    @Override
    public void getTaskProcessTracker(TaskEntity taskEntity, String clientId, String promptId, String orderNum) {
        try {
            taskProcessTracker.getTaskProcess(promptId, orderNum, defaultTaskFinishHandler);
        } catch (IOException e) {
            throw new BizException("获取进度失败");
        }
    }
}
